from __future__ import print_function
from bcc import BPF
import argparse
import os
import json
from event_data_parser import parse_event_data
import time

# arguments
examples = """examples:
    python io_trace.py           # trace all block I/O
    python io_trace.py -d sda    # trace sdc only
"""
parser = argparse.ArgumentParser(
    description="Trace block I/O",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("-d", "--disk", type=str,
    help="trace this disk only")
parser.add_argument("--ebpf", action="store_true",
    help=argparse.SUPPRESS)
parser.add_argument('-o', '--output',
            type=str,
            help='Specify the output JSON file name',
            default='event_data.json')  
parser.add_argument("-t", "--time",
            type=int,
            help="Specify the duration (in seconds) to collect data",
            default=10) 
args = parser.parse_args()
debug = 0

# define BPF program
with open('io_trace.c') as f:
    bpf_text = f.read()
if BPF.kernel_struct_has_field(b'request', b'rq_disk') == 1:
    bpf_text = bpf_text.replace('__RQ_DISK__', 'rq_disk')
else:
    bpf_text = bpf_text.replace('__RQ_DISK__', 'q->disk')

if args.disk is not None:
    disk_path = os.path.join('/dev', args.disk)
    if not os.path.exists(disk_path):
        print("no such disk '%s'" % args.disk)
        exit(1)

    stat_info = os.stat(disk_path)
    dev = os.major(stat_info.st_rdev) << 20 | os.minor(stat_info.st_rdev)

    disk_filter_str = """
    if(key.dev != %s) {
        return 0;
    }
    """ % (dev)

    bpf_text = bpf_text.replace('DISK_FILTER', disk_filter_str)
else:
    bpf_text = bpf_text.replace('DISK_FILTER', '')

if debug or args.ebpf:
    print(bpf_text)
    if args.ebpf:
        exit()

# initialize BPF
b = BPF(text=bpf_text)
if BPF.get_kprobe_functions(b'__generic_make_request_checks'):
    b.attach_kprobe(event="__generic_make_request_checks", fn_name="trace_blkq")
elif BPF.get_kprobe_functions(b'generic_make_request_checks'):
    b.attach_kprobe(event="generic_make_request_checks", fn_name="trace_blkq")

if BPF.tracepoint_exists("block", "block_io_start"):
    b.attach_tracepoint(tp="block:block_io_start", fn_name="trace_blkg")
elif BPF.get_kprobe_functions(b'__blk_account_io_start'):
    b.attach_kprobe(event="__blk_account_io_start", fn_name="trace_blkg")
elif BPF.get_kprobe_functions(b'blk_account_io_start'):
    b.attach_kprobe(event="blk_account_io_start", fn_name="trace_blkg")
else:
    print("ERROR: No found any block io start probe/tp.")
    exit()

if BPF.tracepoint_exists("block", "block_rq_insert"):
    b.attach_tracepoint(tp="block:block_rq_insert", fn_name="trace_blki")

if BPF.get_kprobe_functions(b'__bio_attempt_back_merge'):
    b.attach_kprobe(event="__bio_attempt_back_merge", fn_name="trace_blkm")
elif BPF.get_kprobe_functions(b'bio_attempt_back_merge'):
    b.attach_kprobe(event="bio_attempt_back_merge", fn_name="trace_blkm")


b.attach_kprobe(event="blk_mq_start_request", fn_name="trace_blkd")
if BPF.tracepoint_exists("block", "block_io_done"):
    b.attach_tracepoint(tp="block:block_io_done", fn_name="trace_blkc")
elif BPF.get_kprobe_functions(b'__blk_account_io_done'):
    b.attach_kprobe(event="__blk_account_io_done", fn_name="trace_blkc")
elif BPF.get_kprobe_functions(b'blk_account_io_done'):
    b.attach_kprobe(event="blk_account_io_done", fn_name="trace_blkc")
else:
    print("ERROR: No found any block io done probe/tp.")
    exit()


# header
print("%-3s %-17s %-14s %-7s %-9s %-1s %-19s %-7s" % ("CPU", "TIME(s)", "COMM", "PID",
    "DISK", "T", "SECTOR", "BYTES"), end="")
print("%-1s " % ("EVENT"))



# cache disk major,minor -> diskname
diskstats = "/proc/diskstats"
disklookup = {}
with open(diskstats) as stats:
    for line in stats:
        a = line.split()
        disklookup[a[0] + "," + a[1]] = a[2]

def disk_print(d):
    major = d >> 20
    minor = d & ((1 << 20) - 1)

    disk = str(major) + "," + str(minor)
    if disk in disklookup:
        diskname = disklookup[disk]
    else:
        diskname = "<unknown>"

    return diskname

rwflg = ""
pattern = ""
start_ts = 0
prev_ts = 0
delta = 0

P_SEQUENTIAL = 1
P_RANDOM = 2


# save and load data
def save_event_to_file(event_data, filename=args.output):
    with open(filename, "a") as f:
        json.dump(event_data, f)
        f.write("\n")

def load_events_from_file(filename=args.output):
    events = []
    with open(filename, "r") as f:
        for line in f:
            events.append(json.loads(line))
    return events

def print_saved_events(filename=args.output):
    events = load_events_from_file(filename)
    for event in events:
        print(event)

# process event
def print_event(cpu, data, size):
    event = b["events"].event(data)

    global start_ts
    if start_ts == 0:
        start_ts = event.ts

    if event.rwflag == 1:
        rwflg = "W"
    else:
        rwflg = "R"

    delta = float(event.ts) - start_ts

    disk_name = disk_print(event.dev)


    event_name = event.name.decode('utf-8', 'replace')
    event_data = {
        "cpu": event.cpu,
        "event": event.event.decode('utf-8', 'replace'),
        "pid": event.pid,
        "disk_name": disk_name,
        "rwflg": rwflg,
        "sector": event.sector,
        "len": event.len,
        "ts": event.ts / 1000000,
        "name": event_name,
        "blk_qts": event.blk_qts / 1000000,
        "blk_gts": event.blk_gts / 1000000,
        "blk_its": event.blk_its / 1000000,
        "blk_mts": event.blk_mts / 1000000,
        "blk_dts": event.blk_dts / 1000000,
        "blk_cts": event.blk_cts / 1000000
    }



    print("%-3s %-11.6f %-14.14s %-7s %-9s %-1s %-19s %-7s" % (
        event.cpu, event.ts / 1000000, event.name.decode('utf-8', 'replace'), event.pid,
        disk_name, rwflg, event.sector, event.len), end="")
    print("%-1s " % (event.event.decode('utf-8', 'replace')))

    save_event_to_file(event_data)

# loop with callback to print_event
# b.trace_print()
b["events"].open_perf_buffer(print_event, page_cnt=64)
try:
    duration = args.time  
    start_time = time.time()
    while True:
        b.perf_buffer_poll()
        time_elapsed = time.time() - start_time
        if time_elapsed >= duration:
            raise KeyboardInterrupt
            break
except KeyboardInterrupt:
    parse_event_data(args.output)
    exit()


